#ifndef ASTRA_PRO_HPP
#define ASTRA_PRO_HPP

#include <opencv2/opencv.hpp>
#include <opencv2/highgui/highgui.hpp>
#include "OpenNI.h"
#include "common_msg/TrackedObject.h"

class AstraPro
{
public:
    cv::Mat mColorRawImg;
    cv::Mat mColorDisImg;
    cv::Mat mDepthRawImg;
    cv::Mat mDepthDisImg;
    int mColorFps, mColorRows, mColorCols;
    int mDepthFps, mDepthRows, mDepthCols;

public:
    AstraPro()
    {
        mColorRawImg.create(480, 640, CV_8UC3);
        mColorDisImg.create(480, 640, CV_8UC3);
        mDepthRawImg.create(240, 320, CV_16UC1);
        mDepthDisImg.create(240, 320, CV_8UC1);
    }

    ~AstraPro()
    {
    }

    int open(int cameraId = 0)
    {
        openni::Status rc = openni::OpenNI::initialize();
        if (rc != openni::STATUS_OK)
        {
            std::cout << "Init failed!" << std::endl;
            std::cout << openni::OpenNI::getExtendedError() << std::endl;
            return -1;
        }

        rc = mDevice.open(openni::ANY_DEVICE);
        if (rc != openni::STATUS_OK)
        {
            std::cout << "Couldn't open device!" << std::endl;
            std::cout << openni::OpenNI::getExtendedError() << std::endl;
            return -2;
        }

        rc = mDepthStream.create(mDevice, openni::SENSOR_DEPTH);
        if (rc != openni::STATUS_OK)
        {
            std::cout << "Couldn't create depth stream!" << std::endl;
            std::cout << openni::OpenNI::getExtendedError() << std::endl;
            return -3;
        }

        // or set Mirror=0 in orbbec.ini
        // depth.setMirroringEnabled(false);

        rc = mDepthStream.start();
        if (rc != openni::STATUS_OK)
        {
            std::cout << "Couldn't start the depth stream!" << std::endl;
            std::cout << openni::OpenNI::getExtendedError() << std::endl;
            return -4;
        }

        if (!mCapture.open(0))
        {
            mCapture.open(1);
        }

        if (!mCapture.isOpened())
        {
            std::cout << "Couldn't start the color stream!" << std::endl;
            return -5;
        }

        return 0;
    }

    void close()
    {
        mCapture.release();
        mDevice.close();
        openni::OpenNI::shutdown();
    }

    int getColorImg()
    {
        mCapture >> mColorRawImg;
        if (mColorRawImg.empty())
        {
            std::cout << "get color image failed!" << std::endl;
            return -1;
        }
        mColorDisImg = mColorRawImg.clone();

        return 0;
    }

    int getDepthImg()
    {
        int streamIndex;
        openni::VideoFrameRef frame;

        openni::VideoStream *pStream[] = {&mDepthStream};
        openni::Status rc = openni::OpenNI::waitForAnyStream(pStream, 1, &streamIndex, OPENNI_READ_WAIT_TIMEOUT_MS);
        if (rc != openni::STATUS_OK)
        {
            std::cout << "Wait stream failed! Timeout is " << OPENNI_READ_WAIT_TIMEOUT_MS << " ms" << std::endl;
            std::cout << openni::OpenNI::getExtendedError() << std::endl;
            return 1;
        }

        switch (streamIndex)
        {
        case 0:
            rc = mDepthStream.readFrame(&mDepthFrame);
            break;
        default:
            std::cout << "Wait frame error! Stream index: " << streamIndex << std::endl;
            return 1;
        }

        if (rc != openni::STATUS_OK)
        {
            std::cout << "Read failed!" << std::endl;
            std::cout << openni::OpenNI::getExtendedError() << std::endl;
            return 2;
        }

        if (mDepthFrame.isValid())
        {
            openni::DepthPixel *pDepth = (openni::DepthPixel *)mDepthFrame.getData();
            mDepthRawImg = cv::Mat(mDepthFrame.getHeight(), mDepthFrame.getWidth(), CV_16UC1, (unsigned char *)pDepth);
            mDepthRawImg.convertTo(mDepthDisImg, CV_8UC1, ALPHA);
            // std::cout << "depth_img: " << depth_frame.getWidth() << " x " << depth_frame.getHeight() << std::endl;
        }

        return 0;
    }

    void showInfo()
    {
        getColorImg();
        getDepthImg();
        mColorFps = mCapture.get(cv::CAP_PROP_FPS);
        mColorRows = mColorRawImg.size[0];
        mColorCols = mColorRawImg.size[1];
        std::cout << "ColorImg: " << mColorCols << " x " << mColorRows << " @ " << mColorFps << " Hz" << std::endl;

        openni::VideoMode depthMode = mDepthStream.getVideoMode();
        mDepthFps = depthMode.getFps();
        mDepthRows = depthMode.getResolutionY();
        mDepthCols = depthMode.getResolutionX();
        std::cout << "DepthImg: " << mDepthCols << " x " << mDepthRows << " @ " << mDepthFps << " Hz" << std::endl;
    }

    // +X: right, +Y: down, +Z: forward
void calcObjPosition(common_msg::TrackedObject &obj, cv::Rect2d &rect)
{
	// object center coordinate in depth_img: 320*240 (color_img: 640*480)
	int objCx = (rect.x / 2) + (rect.width / 4);
	int objCy = (rect.y / 2) + (rect.height / 4);

	int yiStart = std::max(0, objCy - 16);
	int yiStop = std::min(239, objCy + 16);
	int xiStart = std::max(0, objCx - 16);
	int xiStop = std::min(319, objCx + 16);
	int zcTemp = 0, zcCnt = 0, zc = 0;
	for (int yi = yiStart; yi < yiStop; yi += 3)
	{
		for (int xi = xiStart; xi < xiStop; xi += 3)
		{
			zcTemp = mDepthRawImg.at<ushort>(yi, xi);
			if (zcTemp != 0)
			{
				zc += zcTemp;
				zcCnt++;
			}
		}
	}
	if (zcCnt > 1)
	{
		zc = zc / zcCnt;
	}
	else
	{
		zc = 0;
	}

	// pixel coordinate to camera coordinate
	// int zc = mDepthRawImg.at<ushort>(obj_c_y, obj_c_x);
	int xc = int((float)zc * mInvFx * ((objCx * 2) - mU0));
	int yc = int((float)zc * mInvFy * ((objCy * 2) - mV0));

    if (zc < 500 || zc > 5000 || abs(xc) > 2400 || abs(yc) > 1800)
    {
        obj.is_missing = true;  // probably missing object
    }

	// obj coordinate (mm)
	obj.position.x = zc;
	obj.position.y = -xc;
	obj.position.z = -yc;
}

private:
    cv::VideoCapture mCapture;
    openni::Device mDevice;
    openni::VideoStream mDepthStream;
    openni::VideoFrameRef mDepthFrame;
    const int OPENNI_READ_WAIT_TIMEOUT_MS = 50;
    const double ALPHA = 255.0 / 8000.0;
    // camera intrinsic parameters
	const float mFx = 570.3422;
	const float mFy = 570.3422;
	const float mInvFx = 1.0 / mFx;
	const float mInvFy = 1.0 / mFy;
	const float mU0 = 319.5;
	const float mV0 = 239.5;
};

#endif